from collections import namedtuple, defaultdict
import gurobipy as gp
from gurobipy import GRB
import math
import matplotlib.pyplot as plt

# ------------这里是数据预处理部分的代码-------------- #
# Customer：使用namedtuple定义一个名为Customer的数据结构，
# 包含四个字段：index（索引）、demand（需求量）、x（x坐标）、y（y坐标）
Customer = namedtuple("Customer", ['index', 'demand', 'x', 'y'])


# 从指定路径读取数据文件
def read(path):
    with open(path, 'r') as file:
        data = file.read()
    lines = data.split('\n')
    # 分割第一行
    # 并分别提取：points_count : 点的数量; vehicle_count : 车辆数量和 vehicle_capacity : 车辆容量。
    parts = lines[0].split(' ')
    points_count, vehicle_count, vehicle_capacity = int(parts[0]), int(parts[1]), int(parts[2])
    # 初始化一个空列表，接着逐行获取所有点的信息，每个点存储为一个Customer实例。
    points_list = []
    for i in range(1, points_count + 1):
        line = lines[i]
        parts = line.split()
        # 包含四个字段：index（索引）、demand（需求量）、x（x坐标）、y（y坐标）
        points_list.append(Customer(i - 1, int(parts[0]), float(parts[1]), float(parts[2])))
    return points_count, vehicle_count, vehicle_capacity, points_list


# ------------这里是数据预处理部分的代码-------------- #


# ------------这里是其它辅助函数或者工具函数的定义-------------- #
# 计算两个点之间的欧几里得距离
def length(point1, point2):
    return math.sqrt((point1.x - point2.x) ** 2 + (point1.y - point2.y) ** 2)


# 绘制求解路径
# 绘制所有的点，连接两点，绘制所有被选中的边。
def draw(tours):
    for sub in tours:
        for i in range(len(sub) - 1):
            plt.plot([points_ls[sub[i]].x, points_ls[sub[i + 1]].x], [points_ls[sub[i]].y, points_ls[sub[i + 1]].y],
                     c='red')
            plt.scatter(points_ls[sub[i]].x, points_ls[sub[i]].y, c='black')
    plt.show()


# ------------这里是其它辅助函数或者工具函数的定义-------------- #

# ------------这里是模型定义-------------- #
# 混合整数规划（MIP）模型定义
def MIP_model(points_count, vehicle_count, vehicle_capacity, points_list):
    """
    (1)已知数据
    points：创建一个包含所有点索引的列表。
    customers：创建一个包含所有客户点（不包括仓库）的索引列表。
    edges：创建一个包含所有可能边的列表（不包括自环）。
    distance：计算并存储每对点之间的距离。
    """
    points = [i for i in range(points_count)]
    customers = [i for i in range(1, points_count)]
    edges = [(i, j) for i in range(points_count) for j in range(points_count) if i != j]
    distance = {(i, j): length(points_list[i], points_list[j]) for i, j in edges}
    """
    (2)决策变量和目标函数
    """
    # 创建一个新的Gurobi模型，命名为CVRP
    model = gp.Model('CVRP')
    # 设定模型求解时间：单位s
    model.Params.TimeLimit = 60
    # 为【每条边】添加一个【二元决策变量】，表示是否选择该边
    select = model.addVars(edges, vtype=GRB.BINARY, name='select')
    # 为每个【客户点】添加一个【连续变量】，表示流量
    flow = model.addVars(customers, vtype=GRB.CONTINUOUS, name='flow')
    # 设置【目标函数】为【最小化】【总距离】。
    # select.prod(distance)：这个方法调用会计算一个线性表达式，
    # 其中每个变量 select[i, j] 被其对应的系数 distance[i, j] 乘以，然后对所有的 (i, j) 进行求和。
    model.setObjective(select.prod(distance), GRB.MINIMIZE)
    """
    (3)约束条件
    gp.quicksum：这是 Gurobi 中的一个函数，用于【快速计算一组变量的线性组合的和】。
    它的主要作用是提高计算效率，【尤其是在处理大规模模型】时。
    """
    # 每个【客户点】一条出边
    model.addConstrs(gp.quicksum(select[i, j] for j in points if j != i) == 1 for i in customers)
    # 每个【客户点】一条入边
    model.addConstrs(gp.quicksum(select[i, j] for i in points if i != j) == 1 for j in customers)
    # 流量顺序约束
    # >>：这是Gurobi中的条件约束符号，表示“如果...则...”的逻辑关系。
    # 即，如果左侧条件为真，则右侧条件必须成立。
    # 这里去除掉了中心仓库点(i=0,j=0)
    model.addConstrs((select[i, j] == 1) >> (flow[i] + points_list[j].demand == flow[j])
                     for i, j in edges if i != 0 and j != 0)
    # 流量大小（范围）约束
    # 流经点i的流量应该大于等于点i本身的需求
    model.addConstrs(points_list[i].demand <= flow[i] for i in customers)
    # 流经点i的流量应该小于等于车辆可以容纳的总容量
    model.addConstrs(flow[i] <= vehicle_capacity for i in customers)
    # 车辆数量约束
    model.addConstr(gp.quicksum(select[0, j] for j in customers) <= vehicle_count)
    model.addConstr(gp.quicksum(select[i, 0] for i in customers) <= vehicle_count)
    """
    (4)求解模型
    """
    model.optimize()
    arcs = [e for e in edges if select[e].x > 0.99]  # 提取解中选择的边。
    print(arcs)
    hash_table = defaultdict(list)  # 创建一个默认值为列表的字典，用于存储路径。
    # 将选择的边存入字典。
    for e in arcs:
        hash_table[e[0]].append(e[1])
    # 初始化一个空列表，用于存储每辆车的路径。
    vehicle_tours = []
    # 构建每辆车的完整路径。
    for k in range(len(hash_table[0])):
        vehicle_tours.append([])
        vehicle_tours[k].append(0)
        vehicle_tours[k].append(hash_table[0][k])
        last = vehicle_tours[k][-1]
        while last != 0:
            vehicle_tours[k].append(hash_table[last][0])
            last = vehicle_tours[k][-1]
    return vehicle_tours


# ------------这里是模型定义-------------- #


# ------------主程序用于流程控制-------------- #
if __name__ == '__main__':
    # 选择不同的测试数据源
    data_path_list = ['./vrp_16_3_0', './vrp_16_3_1', './vrp_16_3_2',
                      './vrp_16_4_1', './vrp_16_4_2',
                      './vrp_16_5_1', './vrp_16_5_2',
                      './vrp_21_4_1', './vrp_5_3_1']
    path = data_path_list[6]
    # 读取数据
    points_cnt, vehicle_cnt, vehicle_cap, points_ls = read(path)
    # 将数据传入模型并求解
    opt_tours = MIP_model(points_cnt, vehicle_cnt, vehicle_cap, points_ls)
    # 根据模型训练的结果绘制求解路径
    draw(opt_tours)
# ------------主程序用于流程控制-------------- #
